Parallel ensemble I/O for State and Increment - #1243
Conversation
Rebased onto develop after the direct-netCDF soca_io_mod PR (#1242) landed. End-to-end ensemble-parallel I/O on top of the stateless direct-netCDF reader. Opt-in via the new oops StateSet/IncrementSet bulk-readEnsemble/writeEnsemble dispatch path; soca provides the model-level implementations. State/Increment: - State::{readEnsemble,writeEnsemble} and Increment::{readEnsemble, writeEnsemble} (Fields + interface Fortran bindings to match) - writeEnsemble: gather each member to a strided writer PE (soca_io_ensemble_root_pe assignment), then phase-2 per-member netCDF writes happen concurrently across writer PEs - readEnsemble: per-member loop today, honoring single-state read mode (strided vs scatter); parallel-across-members reads via soca_io_readers_commit_ensemble are wired for a follow-up soca_io_mod: - writer staging split into define/gather/write/close phases plus soca_io_writers_commit_ensemble; reader staging mirrored with read/distribute/close and soca_io_readers_commit_ensemble. Reader opens+closes ncid inside stage_read (no global file-handle cache, matching develop's stateless model post-2d7ab3cf) - new public knobs: soca_io_ensemble_write_parallel, soca_io_ensemble_read_scatter, soca_io_single_state_read_scatter, soca_io_async_mpi, soca_io_ensemble_root_pe - mpi_pelist_and_comm helper for the writer/reader PE subsets Geometry: - soca_io_config_from_yaml called from soca_geom_init resolves the parallel/sequential, scatter/strided, async-mpi knobs from geometry.io once; values persist module-level for the run Companion to JCSDA-internal/oops setID fix in StateSet::buildFromConfigs (needed so PseudoModel and HTLM dispatch see correct per-member IDs when StateSet takes the HasReadEnsemble path).
0c72b73 to
3b0674f
Compare
…arallel-io-rebased
The companion oops branch now captures the posterior stddev before the in-place ana_pert mutation but prints it in its original position, so the variance-before-mean testref updates are no longer needed. The PR now touches no testrefs (per-member output is byte-identical to the legacy path).
Correctness: - async gather/scatter: give each in-flight MPI_Igatherv/Iscatterv its own rc/disp column (the shared rc(:)/disp(:) were recomputed each loop iteration while a prior nonblocking collective was still using them) - writer async gather: allocate a 1-element dummy recvbuf on non-root PEs instead of passing an unallocated actual argument to MPI_Igatherv Cleanup: - strip profiling instrumentation (OOPS_PERF prints, system_clock blocks, measurement-only mpp_sync barriers between phases) from the bulk read/write orchestrators - remove dead code: 3 unused public getters, the unused pelist out-arg of the comm helper (now current_mpi_comm), unreachable root_pe/reader_pe sentinel loops - assert all writers/readers in a batch share one FMS domain - correct the stale 'follow-up' doc comment on State::readEnsemble
- factor the per-domain atlas->compute-domain copy + writer enqueue out of soca_fields_write_rst and soca_fields_write_ensemble into a shared soca_fields_enqueue_domain helper, so the single-state and bulk paths no longer carry two copies of the mask/fillvalue and enqueue logic - replace the four separate MPI_Allgather calls of the static compute-domain bounds (copy-pasted across the three bulk I/O routines) with one packed Allgather in a gather_compute_domains helper: 12 collectives per call drop to 3, with the rank-ordered layout the Igatherv/Iscatterv paths rely on preserved exactly
geometry.io.'ensemble batch size' caps how many members are built, committed, and freed per bulk read/write pass. Default 0 keeps the single-batch behavior (all members at once, maximum overlap) and is byte-identical to the prior path; a positive N processes the ensemble in waves of at most N members so peak memory is bounded by one wave's writers/readers plus scratch rather than the whole ensemble. The writer/reader root PE is rotated by position within the batch so each wave still spreads its members one-per-PE across the comm. Validated on the 1deg 7-member LETKF: batch 1/2/none give bit-identical analysis (nccmp -d); batch 1 cuts total memory ~27% (21.4 -> 15.6 GB) at the cost of wall time, matching the serial memory floor.
Both async paths posted every (member, var) collective up front, then did one MPI_Waitall, unpacked all, and freed all buffers at the end. That held every request's transient buffer live simultaneously: on the write side the per-root recv buffers (each a full global field), on the read side the per-root send buffers (full global, packed from gbuf). With the whole ensemble in one batch that doubled the global-field memory at peak. Replace the Waitall + unpack-all + free-all with an MPI_Waitany loop: as each gather/scatter completes, free its send buffer, unpack (on the target root) into the gbuf / caller slice, and free its recv buffer immediately. All requests are still posted up front, so the runtime overlap is unchanged; only the transient high-water drops from all requests to the in-flight set. 1deg 7-member LETKF, single batch: total memory 21.38 -> 19.08 GB (-10.8%), max RSS -10.7%, wall unchanged, analysis bit-identical (nccmp -d) and ctests green.
The ensemble write path carried two ways to serialize: the 'ensemble write: sequential' knob (commit each writer fully, one at a time) and, since the batch-size knob landed, batch 1 (commit one member's writers per pass through the orchestrator). They serve the same purpose -- bound peak memory by not holding the whole ensemble's buffers at once -- and batch 1 already matches a serial write's memory profile, so the separate sequential path was redundant A/B scaffolding. Remove ensemble_write_parallel_, its 'ensemble write' YAML key, the public getter, and the fallback branch; soca_fields_write_ensemble now always commits through the orchestrator and relies on 'ensemble batch size' for serialization. The single-writer commit() method stays (still used by the single-state write, geometry, and balance paths). Also drop two dead nx_g/ny_g locals in the write gather. ctests green.
|
I've just tested it with our 0.25 deg run (but with 30 members), nice!: previous run: The memory seems quite a bit larger, I'll look into it later unless @travissluka you have an idea why. |
Yeah, some memory increase is inevitable. Each ensemble member requires the extra buffer for the full state on read/write. It's a trade-off. You can set That said, I know there are still areas for memory optimization (the async MPI calls have duplicate read/write buffers that I would like to get rid of). But, good to see it's on the right track, thanks for testing ! |
shlyaeva
left a comment
There was a problem hiding this comment.
Review solely based on testing and great timing results. I'm not an expert in new or old i/o, and can't review the code changes well. 🎉
Description
Warning
Make sure you check out the matching oops branch (JCSDA-internal/oops#3316)! This soca branch will still compile and run, but it won't go any faster without the oops dispatch hooks.
Replaces the per-member-serial ensemble read/write path with bulk routines that drive concurrent I/O across rotated reader/writer PEs, with optional MPI-async batching for the gather/scatter collectives. Previously the LETKF was doing all ensemble file writes on PE0... ouch. Ensemble member reads were using "strided" read mode (whereby each PE reads its section of a netcdf file), not the worst, but probably not the best.
This PR uses a
oops::StateSet::readEnsemble/writeEnsembleset of methods that let us bypass the per-member loop, doing ensemble read/write in true parallel.YAML configuration
The default settings now should do parallel ensemble I/O in what is hopefully the most efficient way on an HPC. But there are some knobs to turn if you need. All knobs live under
geometry.io.ensemble readscatter|stridedscatterMPI_Scattervto compute groups; strided = each PE reads its compute-domain tile directly.async mpitrue|falsetrueIgatherv/Iscattervacross all (reader/writer, var) pairs so collectives with different roots overlap. Applies to both read and write.ensemble batch sizeN >= 000= process all members in one batch (maximum I/O overlap, peak memory ~ all members resident at once).N > 0commits the ensemble in waves of at mostNmembers to cap peak memory (setN~ number of PEs); batch1matches a serial write's memory footprint.single state readstrided|scatterstridedMemory
Bulk parallel I/O necessarily keeps more in flight than the old serial path: each rotated writer PE stages a member's whole global field (and on read, each reader PE does the same). Two things keep that bounded:
MPI_Waitany) as each gather/scatter finishes, rather than holding every member's buffer until one bigWaitall. On a 1deg 7-member LETKF this cut peak memory ~11% at no batching, for free (I/O overlap unchanged).ensemble batch sizecaps how many members are resident per pass.ensemble batch size: 1brings the footprint down to the serial path's level; the default single batch trades that memory for maximum overlap. Output is bit-identical across all batch sizes.Expected performance impact
Based on tests with a synthetic memory/file-system delay emulating HPC performance... my assistant forecasts an improvement of 1.5–3× on reads, 10–30× on writes... let's see how good that forecast is!
ensemble read: scatteris faster when the number of PEs/nodes is large. But if read times are slower, it might be worth testing putting this back tostrided, which is the currently implemented methodasync mpi: truein theory should go faster when doing ensemble I/O, but the flag is here in case we want to verifysingle state readi'm honestly not sure which would be faster, in theorystridedshould be faster if number of PEs >> the number of lustre OSTs.ensemble batch size: 0should be the fastest ensemble writes, but also uses the most memory. The knob is there if we need to reduce memory.Issue(s) addressed
Dependencies
This PR depends on:
State/IncrementreadEnsemble/writeEnsembledispatch hooks this PR implements)Merge order: oops#3316 first, then this PR.
build-group=https://github.com/JCSDA-internal/oops/pull/3316
Impact
Behavior is identical to the legacy path. Expected wall-clock speedup on HPC; no expected impact on downstream science code.
Checklist
Note
Validated bit-identical to the legacy path on the soca EnKF/LETKF ctests. Still need to stress test on an HPC.
Did I review my code? Of course not, I figure it's correct judging by the cpu/memory usage, runtime, and the fact the output files are bit identical. Who cares about looking at the code??